import * as THREE from "three";

// ---- Tween.ts ----



export class Tween<T extends number | THREE.Vector3 | THREE.Color> {
  times: number[] = [];
  values: T[] = [];
  loop: boolean = false;
  key: keyof Particle;
  enable = true;

  constructor(timeArray: number[], valueArray: T[], key: keyof Particle) {
    this.times = timeArray;
    this.values = valueArray;
    this.key = key;
  }

  lerp(t: number) {
    let i = 0;
    let n = this.times.length;
    while (i < n && t > this.times[i]) i++;
    if (i === 0) return this.values[0];
    if (i === n) return this.values[n - 1];
    let p = (t - this.times[i - 1]) / (this.times[i] - this.times[i - 1]);
    //  这里的问题在于，我只是判断数组第一个元素，如果要让ts知道，我应该把  this.values[i] this.values[i+1]都判断一下。 ts才知道
    if ((this.values[0] as THREE.Vector3).isVector3)
      return (this.values[i - 1] as THREE.Vector3).clone().lerp(this.values[i] as THREE.Vector3, p);
    else return (this.values[i - 1] as number) + p * (this.values[i] as number  -  (this.values[i - 1] as number));
  }

  update(time: number, particles: Particle[] | Set<Particle>) {
    if (!this.key) return;

    for (let particle of particles) {
      if (typeof particle[this.key].copy === 'function') {
        particle[this.key].copy(this.lerp(particle.age));
      } else {
        particle[this.key] = this.lerp(particle.age);
      }
    }
  }
}

// ---- util.ts ----


const {sin,cos,tan, sqrt, min, max, random,PI} = Math;
const sphrical = new THREE.Spherical();
const q1 = new THREE.Quaternion(),zAxes = new THREE.Vector3(0,0,1),v1 = new THREE.Vector3();
/**
 * 锥形空间，角度不超过180
 * @param dir  基准方向 默认为z
 * @param angle  0-1 
 * @returns 
 */
export function  randomRadian(dir: THREE.Vector3,angle: number): THREE.Vector3 {
    q1.setFromUnitVectors(zAxes,dir);
    const phi  = (random()*angle + 0.5)*PI;
    const theta = (random() -0.5 )*angle *PI;
    const r2 = sin(phi);
    const x = r2*sin(theta);
    const z = r2*cos(theta);
    const y = cos(phi);
    v1.set(x,y,z);
    // v1.applyQuaternion(q1);
    return v1.clone();
}
export function randomDirectionInCone(dir: THREE.Vector3, theta: number): THREE.Vector3 {
    // 随机采样一个单位向量，夹角不超过 theta
    const u = random();
    const v = random();
    const phi = 2 * PI * u;
    const cosTheta = cos(PI * theta);
    const z = cosTheta + (1 - cosTheta) * v;
    const sinT = sqrt(1 - z * z);
    const x = sinT * cos(phi);
    const y = sinT * sin(phi);
    // 构造局部坐标系
    const axis = new THREE.Vector3();
    if (Math.abs(dir.x) < 0.99) axis.set(1,0,0); else axis.set(0,1,0);
    const tangent = new THREE.Vector3().crossVectors(dir, axis).normalize();
    const bitangent = new THREE.Vector3().crossVectors(dir, tangent).normalize();
    // 旋转到目标方向
    const local = new THREE.Vector3();
    local.addScaledVector(tangent, x);
    local.addScaledVector(bitangent, y);
    local.addScaledVector(dir.clone().normalize(), z);
    return local.normalize();
}

// 生成粒子时
export function randomRange   ( vMin: number, vMax: number): number {
    return vMin + Math.random() * (vMax - vMin);
}



export  const getSPEMat = (opt: THREE.PointsMaterialParameters) => {
    const mat = new THREE.PointsMaterial(opt);
    const vertParams = `
            attribute float a_Size;
            attribute float a_Angle;
            varying float v_Angle;
            void main() {
            `;
    const vertSize = `
            gl_PointSize = a_Size * size;
            v_Angle = a_Angle; 
            `;
    const fragParams = `
            vec2 rotateUV(vec2 uv, float angle) {
                float s = sin(angle);
                float c = cos(angle);
                mat2 rotationMatrix = mat2(c, -s, s, c);
                return rotationMatrix * (uv -0.5) + 0.5; 
            }
            varying float v_Angle;
            void main() {
            `;
    const fragAngle = `
    #if defined( USE_MAP ) || defined( USE_ALPHAMAP )
        #if defined( USE_POINTS_UV )
            vec2 uv = vUv;
        #else
            vec2 uv = ( uvTransform * vec3( gl_PointCoord.x, 1.0 - gl_PointCoord.y, 1 ) ).xy;
        #endif
        float c = cos(v_Angle);
        float s = sin(v_Angle);
        uv = rotateUV(uv, v_Angle);
        vec2 v2k =   step( vec2(0.,0.),uv)  * step(uv , vec2(1.,1.));
        gl_FragColor.a *= v2k.x * v2k.y;
    #endif
    #ifdef USE_MAP
        diffuseColor *= texture2D( map, uv );
    #endif
    #ifdef USE_ALPHAMAP
        diffuseColor.a *= texture2D( alphaMap, uv ).g;
    #endif
    `;
    mat.onBeforeCompile = (shader) => {
        shader.vertexShader = shader.vertexShader.replace(
            'void main() {', vertParams
        );
        shader.vertexShader = shader.vertexShader.replace(
            'gl_PointSize = size;', vertSize
        );
        shader.fragmentShader = shader.fragmentShader.replace(
            'void main() {', fragParams
        );
        shader.fragmentShader = shader.fragmentShader.replace(
            '#include <map_particle_fragment>', fragAngle
        );
    };
    return mat;
}


// ---- Particle.ts ----


export class Particle {
    position: THREE.Vector3;
    velocity: THREE.Vector3;
    angleVolecity: number;
    angle: number;
    color: THREE.Vector3;
    alpha: number;
    alive: boolean;
    age: number;
    size: number;
    life: number;
    constructor(position: THREE.Vector3, velocity: THREE.Vector3, life: number, angleVolecity = 0) {
        this.position = position.clone();
        this.velocity = velocity.clone();
        this.life = life;
        this.age = 0;
        this.alive = true;
        this.color = new THREE.Vector3(1, 1, 1);
        this.angleVolecity = angleVolecity;
        this.angle = 0;
        this.alpha = 1;
        this.size = 1;
    }

    update(deltaTime: number) {
        this.age += deltaTime;
        this.position.addScaledVector(this.velocity, deltaTime);
        this.angle += this.angleVolecity * deltaTime;

        if (this.isDead()) {
            this.alive = false;
        }
    }

    reset(position: THREE.Vector3, velocity: THREE.Vector3, life: number, angle = 0, size = 1) {
        this.position.copy(position);
        this.velocity.copy(velocity);
        this.life = life;
        this.angle = angle;
        this.size = size;
        this.age = 0;
        this.alive = true;
    }

    isDead() {
        return this.age >= this.life;
    }
}

// ---- Emitter.ts ----






interface EmitterOpt{
    origin: THREE.Vector3, rate: number, life: number, emitterLife: number, volecity: THREE.Vector3, angle: number, size: number, shape?: string, shapeParams?: { [key: string]: number },
    angleVolecity?: number   ;
    randomVelocity?: { 
        direction?: THREE.Vector3;     // 主方向
        coneAngle?: number;            // 椎体角度
        speedMin?: number;             // 最小速度
        speedMax?: number;             // 最大速度
    },
    material?: THREE.PointsMaterial;
    maxCount?: number;
}

const defaultConfig : EmitterOpt = {
    origin: new THREE.Vector3(0, 0, 0),
    rate: 100,
    life: 6,
    emitterLife: 10,
    volecity: new THREE.Vector3(0, 0, 1),
    randomVelocity: { 
        direction: new THREE.Vector3(0,0,1),     // 主方向
        coneAngle: 0,            // 椎体角度 0就是没有锥体
        speedMin: 3,             // 最小速度
        speedMax: 20,             // 最大速度
    },
    angle: 0,
    size: 1,
    shape: 'point',
    shapeParams: {},
    angleVolecity: 0,
    material: getSPEMat({
        color: new THREE.Color(0xffffff),
        size: 1,
        sizeAttenuation: false,
        vertexColors: true,
        depthTest: true,
        depthWrite: false,
        blending: THREE.AdditiveBlending,
        transparent: true,
        alphaTest: 0.001,
    }),
    maxCount: 1000,
}
export class Emitter {
    origin: THREE.Vector3;
    rate: number;
    life: number;
    particles: Set<Particle>;
    emitterLife: number;
    geometry: THREE.BufferGeometry;
    material: THREE.PointsMaterial;
    initialVelocity: THREE.Vector3;
    initialAngleVelocity: number;
    initialAngle: number;
    initialSize: number;
    reLivePool: Set<Particle>;
    maxCount = 1000;
    points: THREE.Points;
    color = new THREE.Color(0xffffff);
    shape: string = 'point';
    shapeParams: { [key: string]: number } = {};
    randomVelocity: {
         direction?: THREE.Vector3; // 主方向
        // 主方向
        coneAngle?: number; // 椎体角度
        // 椎体角度
        speedMin?: number; // 最小速度
        // 最小速度
        speedMax?: number; // 最大速度
    } ={}
    emitAcc= 0; // 每帧发射粒子数 累加器
    constructor(opts:EmitterOpt) {
        const opt = { 
            ...defaultConfig,
            ...opts
        }
        this.particles = new Set();
        this.geometry = new THREE.BufferGeometry();
        this.origin = opt.origin;
        this.rate = opt.rate;
        this.life = opt.life;
        this.emitterLife = opt.emitterLife;
        this.material = opt.material!;

        this.initialVelocity = opt.volecity;
        this.randomVelocity = opt.randomVelocity!;
        this.initialAngleVelocity = opt.angleVolecity!;
        this.initialAngle = opt.angle;
        this.initialSize = opt.size;

        this.reLivePool = new Set();
        this.maxCount =Math.max(opt.maxCount!, opt.life * opt.rate * 2);
        this.initGeometry();
        this.points = new THREE.Points(this.geometry, this.material);
        this.points.frustumCulled = false;
        opt.shapeParams && (this.shapeParams = opt.shapeParams);
        this.shape = opt.shape!;
    }

    initGeometry() {
        const total = this.maxCount;
        const positions = new Float32Array(total * 3);
        const colors = new Float32Array(total * 3);
        const sizes = new Float32Array(total).fill(0);
        const angles = new Float32Array(total);
        let index = 0;
        let particle: Particle;
        for (particle of this.particles) {
            positions[index * 3] = particle.position.x;
            positions[index * 3 + 1] = particle.position.y;
            positions[index * 3 + 2] = particle.position.z;
            colors[index * 3] = particle.color.x;
            colors[index * 3 + 1] = particle.color.y;
            colors[index * 3 + 2] = particle.color.z;
            sizes[index] = particle.size;
            angles[index] = particle.angle;
        }
        this.geometry.setAttribute('position', new THREE.BufferAttribute(positions, 3));
        this.geometry.setAttribute('color', new THREE.BufferAttribute(colors, 3));
        this.geometry.setAttribute('a_Size', new THREE.BufferAttribute(sizes, 1));
        this.geometry.setAttribute('a_Angle', new THREE.BufferAttribute(angles, 1));
    }

    emit(deltaTime: number) {
        if (this.emitterLife < 0) {
            return;
        }

        let count = deltaTime * this.rate;
        if (this.emitterLife < deltaTime) {
            count = this.rate * this.emitterLife;
        }
        this.emitAcc += count;
        //  上次剩余的有小数依然累计 直到满1 再次发射
        count = Math.floor(this.emitAcc);
        this.emitAcc -= count;
        
        let particle: Particle, particlePool = [...this.reLivePool];
        const v1 = new THREE.Vector3();
        for (let index = 0; index < count; index++) {
            v1.set(0,0,0);
            let pos = this.getOrigin(index)!;
            if(this.randomVelocity.coneAngle){
                v1.copy(randomRadian(this.randomVelocity.direction!, this.randomVelocity.coneAngle!));
                v1.multiplyScalar(randomRange(this.randomVelocity.speedMin!, this.randomVelocity.speedMax!));
            }
            v1.add(this.initialVelocity);
            if (particlePool.length > 0) {
                particle = particlePool.pop()!;
                this.reLivePool.delete(particle);
                particle.reset(pos, v1, this.life, this.initialAngle, this.initialSize);
            } else {
                particle = new Particle(pos, v1, this.life, this.initialAngleVelocity);
            }
            this.particles.add(particle);
        }
    }

    getOrigin(index: number) {
        if (this.shape === 'point') return this.origin;
        if (this.shape === 'box') {
            let x = (Math.random() - 0.5) * this.shapeParams.x;
            let y = (Math.random() - 0.5) * this.shapeParams.y;
            let z = (Math.random() - 0.5) * this.shapeParams.z;
            return new THREE.Vector3(x, y, z).add(this.origin);
        }
    }

    update(deltaTime: number) {
        this.emitterLife -= deltaTime;
        this.emit(deltaTime);
        let particles = [...this.particles];
        for (let particle of Array.from(particles)) {
            particle.update(deltaTime);
            if (!particle.alive) {
                this.particles.delete(particle);
                this.reLivePool.add(particle);
            }
        }
        this.updateGeometry();
    }

    updateGeometry() {
        const total = this.particles.size;
        if (total <= this.maxCount) {
            let positions = this.geometry.attributes.position.array as Float32Array;
            let colors = this.geometry.attributes.color.array as Float32Array;
            let sizes = this.geometry.attributes.a_Size.array as Float32Array;
            let angles = this.geometry.attributes.a_Angle.array as Float32Array;
            let particle: Particle, particles = [...this.particles];
            for (let index = 0; index < total; index++) {
                particle = particles[index];
                positions[index * 3] = particle.position.x;
                positions[index * 3 + 1] = particle.position.y;
                positions[index * 3] = particle.position.x;
                positions[index * 3 + 1] = particle.position.y;
                positions[index * 3 + 2] = particle.position.z;
                colors[index * 3] = particle.color.x;
                colors[index * 3 + 1] = particle.color.y;
                colors[index * 3 + 2] = particle.color.z;
                sizes[index] = particle.size;
                angles[index] = particle.angle;
            }
            this.geometry.setDrawRange(0, total);
            this.geometry.attributes.position.needsUpdate = true;
            this.geometry.attributes.color.needsUpdate = true;
            this.geometry.attributes.a_Size.needsUpdate = true;
            this.geometry.attributes.a_Angle.needsUpdate = true;
        } else {
            // 需要重新生成geometry todo
        }
    }

    reLive(emitterLife: number) {
        this.emitterLife = emitterLife;
    }
}



// ---- ForceField.ts ----



export class ForceField {
    enabled = true;
    getForce(particle: Particle) {
        return new THREE.Vector3(0, 0, 0);
    }
    update(deltaTime: number, particles: Set<Particle> | Particle[]) {
        for (let particle of particles) {
            let force = this.getForce(particle);
            particle.velocity.addScaledVector(force, deltaTime);
        }
    }
}

export class GravityForceField extends ForceField {
    gravity: THREE.Vector3;
    constructor(gravity: THREE.Vector3) {
        super();
        this.gravity = gravity;
    }
    getForce() {
        return this.gravity;
    }
}

// ---- Constaint.ts ----



/**
 * 约束接口，所有约束需实现 apply 方法
 */
export interface IConstraint {
    enabled: boolean;
    /**
     * 对单个粒子应用约束
     */
    apply(particle: Particle): void;
    update( particles: Iterable<Particle>): void;
}

/**
 * 示例：地面碰撞约束（y=0为地面）
 */
export class GroundConstraint implements  IConstraint {
    enabled = true;
    restitution: number; // 反弹系数
    baseHeight= 0; // 地面高度
    constructor(restitution = 0.8,baseHeight = 0) {
        this.restitution = restitution;
        this.baseHeight = baseHeight;
    }

    apply(particle: Particle): void {
        if (particle.position.y < this.baseHeight) {
            particle.position.y = this.baseHeight;
            if (particle.velocity.y < 0) {
                particle.velocity.y *= -this.restitution;
            }
        }
    }
    update(particles: Iterable<Particle>): void {
        for (let particle of particles) {
            this.apply(particle);
        }
    }
}


/**
 * 通用平面约束
 * @param planeNormal 平面法向量（会自动归一化）
 * @param planeConstant 平面常数（平面方程 n·x + d = 0）
 * @param restitution 反弹系数（能量损失，0~1，1为完全弹性）
 * @param friction 摩擦系数（0~1，1为完全摩擦，0为无摩擦）
 */
export class PlaneConstraint implements IConstraint {
    enabled = true;
    normal: THREE.Vector3;
    constant: number;
    restitution: number;
    friction: number;

    constructor(
        planeNormal: THREE.Vector3,
        planeConstant: number = 0,
        restitution: number = 0.8,
        friction: number = 0.2
    ) {
        this.normal = planeNormal.clone().normalize();
        this.constant = planeConstant;
        this.restitution = restitution;
        this.friction = friction;
    }

    apply(particle: Particle): void {
        // 点到平面距离
        const dist = this.normal.dot(particle.position) + this.constant;
        if (dist < 0) {
            // 修正位置到平面上
            particle.position.addScaledVector(this.normal, -dist);

            // 速度分解
            const v = particle.velocity;
            const vN = this.normal.clone().multiplyScalar(v.dot(this.normal)); // 垂直分量
            const vT = v.clone().sub(vN); // 切向分量

            // 垂直分量反弹（能量损失）
            vN.multiplyScalar(-this.restitution);

            // 切向分量摩擦
            vT.multiplyScalar(1 - this.friction);

            // 合成新速度
            particle.velocity.copy(vN.add(vT));
        }
    }
    update( particles: Iterable<Particle>): void {
        for (let particle of particles) {
            this.apply(particle);
        }
    }
}

// ---- ParticleSystem.ts ----





export class ParticleSystem {
    emitters: Emitter[];
    forces: ForceField[];
    constraints: IConstraint[] = [];
    tweens: Tween<any>[] = [];
    time: number = 0;

    constructor(emitters: Emitter[], forces: ForceField[] = [], tweens: Tween<any>[] = [],constraints:IConstraint[]=[]) {
        this.emitters = emitters;
        this.forces = forces;
        this.tweens = tweens;
        this.constraints = constraints;
    }

    update(deltaTime: number) {
        this.time += deltaTime;
        for (let emitter of this.emitters) {
            for (let force of this.forces) {
                force.enabled && force.update(deltaTime, emitter.particles);
            }
            for (let constraint of this.constraints) {
                constraint.enabled && constraint.update( emitter.particles);
            }
            for (let tween of this.tweens) {
                tween.enable&& tween.update(this.time, emitter.particles);
            }
            emitter.update(deltaTime);
        }
    }
}